import numpy as np
from matplotlib import pyplot as plt
import types
from time import sleep
import h5py
import pickle
import subprocess
import pdb
"""
Module containing class to simulate agitated filter dryer, and a classes to 
help set-up and run the calculation.
"""

class Sim():
    """
    Class to solve the dynamic drying in an agitated filter dryer
    Model is based on Stepanek paper (Murro et al, Model-based scale-up of vacuum contact
    drying of pharmaceutical compounds, Chemical Engineering Science, 66
    5045-5054 2011)
    Too run model
        sim = AFD()
        sim.parameters.update(new_parameters)
        sim.set_mesh_size(*mesh_size)
        sim.run()
    """

    def __init__(self):
        """
        Instantiation
        """
        self.t = []
        self.v_L = []
        self.T = []
        self.parameters = {
                'R'     : 0.05,
                'hq'    : 200,
                'H'     : 0.04,
                'w_Lo'   : 0.31,
                'To'    : 323.,
                'pG'    : 4000.,
                'Th'    : 323.,
                'H_vap' : 682.,
                'A'     : 8.37895,
                'B'     : 1788.020,
                'C'     : 227.438,
                'alpha' : 5e-7,
                'v_Lo'   : 0.001,
                'd_S'    : 2.2e3,
                'd_L'    : 804,
                'Cp_S'   : 837,
                'Cp_L'   : 2371,
                'l_L'   : 0.1455,
                'l_D'   : 0.058,
                'v_S'    : 0.256,
                'k_vap' : 1e-4,
                'ag_t'  : 1e7,
                'dt'    : 1e-3,
                'sim_time' : 15000
                }
        self.parameter_definitions = {
                'R'       : 'dryer radius (m)',
                'hq'      : 'wall heat transfer (W/m^2K)',
                'H'       : 'cake height (m)',
                'w_Lo'     : 'initial moisture content ( (mass wet - mass dry )/ mass wet bed )',
                'To'      : 'initial temperature (K)',
                'pG'      : 'dryer headspace pressure (Pa)',
                'Th'      : 'jacket temperature (K)',
                'H_vap'   : 'latent heat of vaporisation (kJ/kg)',
                'A'       : 'Antoine constant (parameters given in C and mmHg)',
                'B'       : 'Antoine constant (parameters given in C and mmHg)',
                'C'       : 'Antoine constant (parameters given in C and mmHg)',
                'alpha'   : 'parameter for diffusion coefficient',
                'v_Lo'     : 'parameter for diffusion coefficient',
                'd_S'     : 'solid density (kg/m^3)',
                'd_L'     : 'liquid density (kg/m^3)',
                'Cp_S'     : 'solid heat capacity (J/kg K)',
                'Cp_L'     : 'liquid heat capacity (J/kg K)',
                'l_L'     : 'thermal conductivity liquid (W/mK)',
                'l_D'     : 'effective thermal conductivity dry (W/mK)',
                'v_S'      : 'solid volume fraction',
                'k_vap'   : 'vaporisation constant (s)',
                'ag_t'    : 'agitation interval (s)',
                'dt'      : 'time increment (s)',
                'sim_time': 'total simulation time (s)'
                }
        self.mesh = {
                'Nr':15,
                'Nz':20,
                'Nt':10000
                }
        self.create_mesh()

    def set_parameters(self,new_parameters):
        """
        Set the AFD parameters
        new_parameters - dictionary of new parameter values
        """
        for k,v in new_parameters.items():
            if not k in set(self.parameters):
                raise ValueError("%s not in parameters"%k)
        self.parameters.update(new_parameters)

    def set_mesh_size(self,Nr,Nz,Nt):
        """
        set mesh size
        Nr  - number of radial nodal points
        Nz  - number of axial nodal points
        Nt  - number of time nodal points
        """
        self.mesh.update({'Nr':Nr,'Nz':Nz,'Nt':Nt})

    def create_mesh(self):
        """
        Create arrays containing the nodal points
        """
        R = self.parameters['R']
        H = self.parameters['H']
        Nr = self.mesh['Nr']
        Nz = self.mesh['Nz']
        Nt = self.mesh['Nt']
        self.mesh.update({'dr':R/float(Nr-1)})
        self.mesh.update({'dz':H/float(Nz-1)})
        # start first radial node just short of zero to avoid singularities (e.g. r=0 symmetry)
        [self.mesh['z'],self.mesh['r']] = np.meshgrid(np.linspace(0,H,Nz),np.linspace(1e-2*self.mesh['dr'],R,Nr))
        # data variables
        self.t = np.zeros((Nt))
        self.v_L = np.zeros((Nt,Nr,Nz))
        self.T = np.zeros((Nt,Nr,Nz))

    def calc_initial_conditions(self):
        """
        Set the initial conditions of the nodes
        """
        w_Lo = self.parameters['w_Lo']
        To = self.parameters['To']
        # Set first time node point
        self.T[0] = np.zeros(self.T.shape[1:]) + To
        self.v_L[0] = np.zeros(self.v_L.shape[1:]) + self.calc_vol_fraction(w_Lo)

    def increment(self,v_L,T):
        """
        Calculate the incremental change in the dependent variables
        v_L - current liquid volume fraction 
        T - current temperature
        return [dv,dT] differential liquid volume fraction and differential Temperature change
        """
        r = self.mesh['r']
        z = self.mesh['z']
        d_L = self.parameters['d_L']
        H_vap = self.parameters['H_vap'] * 1e3  # change kJ -> J units
        # evaporation rate
        pv = self.vap_press(T)
        mLG = self.evap_rate(pv,v_L)
        # average heat capacity
        d_Cp_avg = self.avg_r_Cp(v_L,T)
        # liquid diffusivity
        DL = self.diffusivity(v_L)
        # effective conductivity
        l_eff = self.eff_cond(v_L)
        # partial derivatives
        [dvdr,d2vdr2,dvdz,d2vdz2,dTdr,d2Tdr2,dTdz,d2Tdz2,dDLdr,dDLdz,dldr,dldz] = self.partials(v_L,T,DL,l_eff)
        # increment
        dv = DL * ( d2vdr2 + 1./r*dvdr + d2vdz2 ) + dvdr * dDLdr + dvdz * dDLdz - mLG / d_L
        dT = l_eff / d_Cp_avg * ( (d2Tdr2 + 1./r*dTdr + d2Tdz2) + dTdr * dldr + dTdz * dldz - mLG * H_vap )
        return [dv,dT]

    def partials(self,v_L,Te,DL,l):
        """
        Calculate the partial derivatives
        v_L - liquid volume fraction
        Te - Temperature
        DL - liquid diffusivity
        l - effective conductivity
        return list of partial derivatives
        """
        dr = self.mesh['dr']
        dz = self.mesh['dz']
        hq = self.parameters['hq']
        Th = self.parameters['Th']
        # create arrays with phantom nodes
        v = np.zeros((v_L.shape[0]+2,v_L.shape[1]+2))
        T = np.zeros((Te.shape[0]+2,Te.shape[1]+2))
        v[1:-1,1:-1] = v_L
        T[1:-1,1:-1] = Te
        # set phantom node values by boundary conditions
        # r=R (0 mass flux, equal thermal wall flux)
        v[-1,1:-1] = v[-3,1:-1]
        T[-1,1:-1] = 2 * dr * hq * ( Th - T[-2,1:-1] ) / l[-1,:] + T[-3,1:-1]
        # z=0 (0 mass flux, equal thermal wall flux)
        v[1:-1,0] = v[1:-1,2]
        T[1:-1,0] = 2 * dz * hq * (Th - T[1:-1,1] ) / l[:,0] + T[1:-1,2]
        # z=H (0 mass flux and thermal flux at surface)
        v[1:-1,-1] = v[1:-1,-3]
        T[1:-1,-1] = T[1:-1,-3]
        # r=0 (symmetry)
        v[0,1:-1] = v[2,1:-1]
        T[0,1:-1] = T[2,1:-1]
        # corners
        T[-1,0] = np.average([T[-2,0],T[-1,1]])
        # partial derivatives
        dvdr = ( v[2:,1:-1] - v[1:-1,1:-1] ) / dr
        dTdr = ( T[2:,1:-1] - T[1:-1,1:-1] ) / dr
        d2vdr2 = ( v[2:,1:-1] - 2*v[1:-1,1:-1] + v[:-2,1:-1] ) / dr**2
        d2Tdr2 = ( T[2:,1:-1] - 2*T[1:-1,1:-1] + T[:-2,1:-1] ) / dr**2
        dvdz = ( v[1:-1,2:] - v[1:-1,1:-1] ) / dz
        dTdz = ( T[1:-1,2:] - T[1:-1,1:-1] ) / dz
        d2vdz2 = ( v[1:-1,2:] - 2*v[1:-1,1:-1] + v[1:-1,:-2] ) / dz**2
        d2Tdz2 = ( T[1:-1,2:] - 2*T[1:-1,1:-1] + T[1:-1,:-2] ) / dz**2
        dDLdr = np.zeros(v_L.shape)
        dDLdr[:-1,:] = ( DL[1:,:] - DL[:-1,:]) / dr
        dDLdr[-1,:] = dDLdr[-2,:]
        dDLdz = np.zeros(v_L.shape)
        dDLdz[:,:-1] = ( DL[:,1:] - DL[:,:-1] ) / dz
        dDLdz[:,-1] = dDLdz[:,-2]
        dldr = np.zeros(v_L.shape)
        dldr[:-1,:] = ( l[1:,:] - l[:-1,:] ) / dr
        dldr[-1,:] = dldr[-2,:]
        dldz = np.zeros(v_L.shape)
        dldz[:,:-1] = ( l[:,1:] - l[:,:-1] ) / dz
        dldz[:,-1] = dldz[:,-2]
        return [dvdr,d2vdr2,dvdz,d2vdz2,dTdr,d2Tdr2,dTdz,d2Tdz2,dDLdr,dDLdz,dldr,dldz]

    def diffusivity(self,v_L):
        """
        Calculate the diffusivity at all nodal points
        v_L - liquid volume fraction
        return - liquid diffusivity
        """
        alpha = self.parameters['alpha']
        v_Lo = self.parameters['v_Lo']
        DL = alpha * (v_L - v_Lo)
        return DL

    def vap_press(self,T):
        """
        Calculate the vapor pressure at all nodal points
        T - temperature (K)
        return vaport pressure (Pa)
        """
        A = self.parameters['A']
        B = self.parameters['B']
        C = self.parameters['C']
        T = T - 273.
        pv = 10**(A - B/(C+T))
        # change mmHg -> Pa units
        pv = pv * 1.01325e5 / 760.
        return pv

    def evap_rate(self,pv,v_L):
        """
        Calculate the evaporation rate at all nodal points
        pv - vapor pressure
        v_L - liquid volume fraction
        return - liquid evaporation rate
        """
        k_vap = self.parameters['k_vap']
        pG = self.parameters['pG']
        d_L = self.parameters['d_L']
        mLG = k_vap * d_L * (pv - pG) / pG
        mLG[v_L<=0] = 0
        mLG[pv<=pG] = 0
        if np.min(mLG<0):
            print "mLG<0"
            pdb.set_trace()
        return mLG

    def avg_r_Cp(self,v_L,T):
        """
        Calculate the average density times heat capacity (heat capacity per unit volume) at all nodal points
        v_L - liquid volume fraction
        T - Temperature
        return bulk average heat capacity
        """
        d_L = self.parameters['d_L']
        d_S = self.parameters['d_S']
        Cp_L = self.parameters['Cp_L']
        Cp_S = self.parameters['Cp_S']
        v_S = self.parameters['v_S']
        r_Cp_avg = v_L*d_L*Cp_L + v_S*d_S*Cp_S
        return r_Cp_avg

    def gas_density(self,T):
        """
        Calculate the gas density
        T - Temperature
        return gas density
        """
        R = 8.314
        pG = self.parameters['pG']
        density = pG/(R*T) * 30e-3
        return density

    def eff_cond(self,v_L):
        """
        Calculate the effective thermal conductivity at all nodal points
        v_L - liquid volume fraction
        return effective conductivity
        """
        l_wet = self.parameters['l_L']
        l_dry = self.parameters['l_D']
        v_S = self.parameters['v_S']
        l_eff = l_dry + v_L/(1-v_S) * (l_wet - l_dry)
        return l_eff

    def wt_fraction(self,v_L):
        """
        Calculate the weight fraction from the volume fraction
        """
        d_L = self.parameters['d_L']
        d_S = self.parameters['d_S']
        v_S = self.parameters['v_S']
        wL = v_L*d_L / (v_S*d_S + v_L*d_L)
        return wL

    def calc_vol_fraction(self,wL):
        """
        Calculate the volume fraction from the weight fraction
        """
        d_L = self.parameters['d_L']
        d_S = self.parameters['d_S']
        v_S = self.parameters['v_S'] 
        v_L = v_S*(d_S/d_L)*(wL/(1-wL))
        return v_L

    def agitate(self,v_L,T):
        """
        Calculate the end temperature and liquid volume fraction after the
        cake has been spatially homoginized
        v_L - liquid volume fraction
        T - temperature
        return [v_L_h,T_h] homoginized liquid volume fraction and temperature
        """
        d_Cp_avg = self.avg_r_Cp(v_L,T)
        v_L_h = np.zeros(v_L.shape) + np.average(v_L)
        T_h = np.zeros(T.shape) + np.average(d_Cp_avg * T)/np.average(d_Cp_avg)
        return [v_L_h,T_h]

    def run(self):
        """
        Perform the calculation
        """
        self.create_mesh()
        self.calc_initial_conditions()
        num_loops = np.int(self.parameters['sim_time'] / self.parameters['dt'])
        save_freq = np.int(num_loops / self.v_L.shape[0])
        dt = self.parameters['dt']
        t = 0
        v_L = self.v_L[0]
        T = self.T[0] 
        interval_time = 0
        for i in xrange(1,num_loops):
            [dv_L,dT] = self.increment(v_L,T)
            t = t + dt
            v_L = v_L + dv_L * dt
            T = T + dT * dt
            v_L[v_L<0] = 0
            interval_time = interval_time + dt
            if interval_time > self.parameters['ag_t']:
                [v_L,T] = self.agitate(v_L,T)
                interval_time = 0
            if np.mod(i,save_freq) == 0:
                node = np.int(i / save_freq)
                self.t[node] = t
                self.v_L[node] = v_L
                self.T[node] = T

class Batch():
    """
    Class to run batches of the simulation
    """
    def __init__(self,fname,runs):
        """
        fname - string name of hdf5 file to save data
        runs - list of dictionaries ['name':'name of run','description':'description of run','parameters':{parameters},'mesh_size':[Nr,Nz,Nt]]
        """
        self.fname = fname
        self.runs = runs

    def save_run(self,f_h5py,group_name,run_name,description,sim):
        """
        Save the simulation data in an hdf5 file
        fname - string name of hdf5 file
        run_name - string name for group in hdf5 file to instert simulation data
        description - string to add as attribute to group in hdf5 file
        sim - instance of afd.AFD()
        """
        f_h5py = h5py.File('%s'%(self.fname))
        try:
            f_h5py.create_group(group_name)
        except ValueError:
            pass
        full_name = group_name+'/'+run_name
        f_h5py.create_group(full_name)
        f_h5py[full_name].attrs['description'] = description
        f_h5py.create_dataset('%s/parameters'%full_name,data=np.array(sim.parameters.items()))
        f_h5py.create_dataset('%s/t'%full_name,data=sim.t)
        f_h5py.create_dataset('%s/v_L'%full_name,data=sim.v_L)
        f_h5py.create_dataset('%s/T'%full_name,data=sim.T)
        f_h5py.create_dataset('%s/r'%full_name,data=sim.mesh['r'])
        f_h5py.create_dataset('%s/z'%full_name,data=sim.mesh['z'])
        #pickle_sim = pickle.dumps(sim)
        #f_h5py.create_dataset('%s/sim'%run_name,data=pickle_sim)
        f_h5py.close()

    def save_plots(self,run_name,sim):
        an = Analysis()
        plots = Plots()
        plots.animated_plots(sim.mesh['r'],sim.mesh['z'],sim.v_L,sim.T,run_name)
        fig_temp = plots.temperature(sim.t,sim.T)
        fig_temp.savefig('%s_temperature.png'%run_name)
        w_L = an.mass_fraction(sim.v_L,sim.parameters['v_S'],sim.parameters['d_L'],sim.parameters['d_S'])
        fig_conc = plots.solvent_conc(sim.t,w_L)
        fig_conc.savefig('%s_concentration.png'%run_name)

    def run_batch(self):
        """
        Run the AFD simulation with different parameters
        """
        for run in self.runs:
            group_name = run['group_name']
            run_name = run['run_name']
            description = run['description']
            new_parameters = run['parameters']
            mesh_size = run['mesh_size']
            sim = AFD()
            sim.parameters.update(new_parameters)
            sim.set_mesh_size(*mesh_size)
            sim.run()
            self.save_run(self.fname,group_name,run_name,description,sim)
            self.save_plots(run_name,sim)

if __name__=="__main__":
    """
    syntax: python afd.py fname
    fname - string for hdf5 file name
    """
    import sys
    import afd
    fname = sys.argv[1]
    runs = [
            {
            'group_name':'mesh',
            'run_name':'20_30_10000',
            'description':'vary mesh size',
            'parameters':{'To':323,'hq':200,'v_S':0.256,'k_vap':1e-3,'sim_time':5000},
            'mesh_size':[20,30,10000],
            },
            {
            'group_name':'mesh',
            'run_name':'40_60_10000',
            'description':'vary mesh size',
            'parameters':{'To':323,'hq':200,'v_S':0.256,'k_vap':1e-3,'sim_time':5000},
            'mesh_size':[40,60,10000],
            },
            {
            'group_name':'mesh',
            'run_name':'80_120_10000',
            'description':'vary mesh size',
            'parameters':{'To':323,'hq':200,'v_S':0.256,'k_vap':1e-3,'sim_time':5000},
            'mesh_size':[80,120,10000],
            },
            {
            'group_name':'mesh',
            'run_name':'40_60_20000',
            'description':'vary mesh size',
            'parameters':{'To':323,'hq':200,'v_S':0.256,'k_vap':1e-3,'sim_time':5000},
            'mesh_size':[40,60,20000],
            },
            {
            'group_name':'mesh',
            'run_name':'40_60_40000',
            'description':'vary mesh size',
            'parameters':{'To':323,'hq':200,'v_S':0.256,'k_vap':1e-3,'sim_time':5000},
            'mesh_size':[40,60,40000],
            }
           ]
    batch = afd.Batch(fname,runs)
    batch.run_batch()
